# DQN for CartPole in OpenAI Gym

# https://zhuanlan.zhihu.com/p/21477488

# tensorboard --logdir logs

import gym
import tensorflow as tf 
import numpy as np 
import random
from collections import deque # 容器数据类型:deque类似列表(list)的容器，实现了在两端快速添加(append)和弹出(pop)
import sys
sys.path.append('../highway-env')
import highway_env

# Hyper Parameters for DQN
GAMMA = 0.9 			# discount factor for target Q 
INITIAL_EPSILON = 0.5 	# starting value of epsilon
FINAL_EPSILON = 0.01 	# final value of epsilon
REPLAY_SIZE = 100 		# experience replay buffer size
BATCH_SIZE = 32			# size of minibatch

# DQN Agent
class DQN():
	#初始化
	def __init__(self, env): 
		# init experience replay
        # 我们的DQN一个很重要的功能就是要能存储数据，然后在训练的时候minibatch出来。
		# 所以，我们需要构造一个存储机制。这里使用deque来实现
		self.replay_buffer = deque()
		# init some parameters
		self.time_step = 0
		self.epsilon = INITIAL_EPSILON
		# self.state_dim = env.observation_space.shape[0]  # (5,5); 5
		self.state_dim = 25
		self.action_dim = env.action_space.n  # 5
		print(" env.observation_space.shape ", env.observation_space.shape)
		print(" state_dim ",self.state_dim  ) 
		print(" action_dim ",  self.action_dim )
		
		# create networks
		self.create_Q_network()
		self.create_training_method()
		
		# init session
		self.session = tf.InteractiveSession()
		self.session.run(tf.initialize_all_variables())

		# loading networks
		self.saver = tf.train.Saver()
		checkpoint = tf.train.get_checkpoint_state("saved_networks")
		if checkpoint and checkpoint.model_checkpoint_path:
				self.saver.restore(self.session, checkpoint.model_checkpoint_path)
				print ("Successfully loaded:", checkpoint.model_checkpoint_path)
		else:
				print ("Could not find old network weights")

		# recode the training process
		# tensorboard --logdir logs
		global summary_writer
		summary_writer = tf.summary.FileWriter('logs/',graph=self.session.graph)


	def create_Q_network(self): #创建Q网络
        # 创建最基本的MLP，中间层设置为20
		# network weights
        # 只有一个隐层，然后使用relu非线性单元。相信对MLP有了解的知友看上面的代码很easy		
		W1 = self.weight_variable([self.state_dim,20])
		b1 = self.bias_variable([20])
		W2 = self.weight_variable([20,self.action_dim])
		b2 = self.bias_variable([self.action_dim])
		# input layer
		# [None, self.state_dim]表示列是4，行不定
		self.state_input = tf.placeholder("float",[None,self.state_dim]) # 要注意的是我们state 输入的格式，因为使用minibatch，所以格式是[None,state_dim]
		# hidden layers
		h_layer = tf.nn.relu(tf.matmul(self.state_input,W1) + b1)
		# Q Value layer
		self.Q_value = tf.matmul(h_layer,W2) + b2

	# 创建训练方法
    # 这里的y_input就是target Q值。我们这里采用Adam优化器，其实随便选择一个必然SGD，RMSProp都是可以的.
	# 可能比较不好理解的就是Q值的计算。这里大家记住动作输入是one_hot_key的形式，
	# 因此将Q_value和action_input向量相乘得到的就是这个动作对应的Q_value,然后用reduce_sum将数据维度压成一维.
	def create_training_method(self): 
		self.action_input = tf.placeholder("float",[None,self.action_dim]) # one hot presentation
		self.y_input = tf.placeholder("float",[None])
		Q_action = tf.reduce_sum(tf.multiply(self.Q_value,self.action_input),reduction_indices = 1)
		self.cost = tf.reduce_mean(tf.square(self.y_input - Q_action))
		tf.summary.scalar("loss",self.cost)
		global merged_summary_op
		merged_summary_op = tf.summary.merge_all()
		self.optimizer = tf.train.AdamOptimizer(0.0001).minimize(self.cost)

    # 感知存储信息
    # 这里需要注意的一点就是动作格式的转换。我们在神经网络中使用的是one hot key的形式，而在OpenAI Gym中则使用单值.
    # 比如我们输出动作是1，那么对应的one hot形式就是[0,1]，如果输出动作是0，那么one hot 形式就是[1,0],这样做的目的是为了之后更好的进行计算.
    # 在perceive中一个最主要的事情就是存储。然后根据情况进行train,这里我们要求只要存储的数据大于Batch的大小就开始训练.
	def perceive(self,state,action,reward,next_state,done):
		one_hot_action = np.zeros(self.action_dim)   # [0,0]
		one_hot_action[action] = 1
		'''
        state   [-1.42546837 -0.7781473   0.13930813  0.60065179]
		one_hot_action  [1. 0.]
		reward  1.0
		next_state   [-1.44103131 -0.9749149   0.15132117  0.9337681 ]
		done  False
		'''
		self.replay_buffer.append((state,one_hot_action,reward,next_state,done))
		if len(self.replay_buffer) > REPLAY_SIZE:
			self.replay_buffer.popleft()
		if len(self.replay_buffer) > BATCH_SIZE:
			self.train_Q_network()

    # 训练网络
    # 首先就是进行minibatch的工作，然后根据batch计算y_batch,最后就是用optimizer进行优化.
	def train_Q_network(self):
		self.time_step += 1
		# Step 1: obtain random minibatch from replay memory 
		# (state,one_hot_action,reward,next_state,done)
		minibatch = random.sample(self.replay_buffer,BATCH_SIZE)
		state_batch = [data[0] for data in minibatch]
		action_batch = [data[1] for data in minibatch]
		reward_batch = [data[2] for data in minibatch]
		next_state_batch = [data[3] for data in minibatch]

		# Step 2: calculate y
		y_batch = []
		Q_value_batch = self.Q_value.eval(feed_dict={self.state_input:next_state_batch})
		for i in range(0,BATCH_SIZE):
			done = minibatch[i][4]
			if done:
				y_batch.append(reward_batch[i])
			else :
				y_batch.append(reward_batch[i] + GAMMA * np.max(Q_value_batch[i]))

		# training
		self.optimizer.run(feed_dict={
			self.y_input:y_batch,
			self.action_input:action_batch,
			self.state_input:state_batch
			})

		summary_str = self.session.run(merged_summary_op,feed_dict={
				self.y_input : y_batch,
				self.action_input : action_batch,
				self.state_input : state_batch
				})
		summary_writer.add_summary(summary_str,self.time_step)

		# save network every 1000 iteration
		if self.time_step % 1000 == 0:
			self.saver.save(self.session, 'saved_networks/' + 'network' + '-dqn', global_step = self.time_step)

    # 输出带随机的动作
    # 区别之前已经说过，一个是根据情况输出随机动作，一个是根据神经网络输出。由于神经网络输出的是每一个动作的Q值，因此我们选择最大的那个Q值对应的动作输出
	def egreedy_action(self,state):
		Q_value = self.Q_value.eval(feed_dict = {
			self.state_input:[state]
			})[0]
		if random.random() <= self.epsilon:
			return random.randint(0,self.action_dim - 1)
		else:
			return np.argmax(Q_value)
        # egreedy的epsilon是不断变小的，也就是随机性不断变小。怎么理解呢？就是一开始需要更多的探索，所以动作偏随机，慢慢的我们需要动作能够有效，因此减少随机
		self.epsilon -= (INITIAL_EPSILON - FINAL_EPSILON)/10000

    #输出确定动作
	def action(self,state):
		return np.argmax(self.Q_value.eval(feed_dict = {
			self.state_input:[state]
			})[0])

	def weight_variable(self,shape):
		initial = tf.truncated_normal(shape)
		return tf.Variable(initial)

	def bias_variable(self,shape):
		initial = tf.constant(0.01, shape = shape)
		return tf.Variable(initial)

# Hyper Parameters
ENV_NAME = 'highway-v0'
EPISODE = 10000 # Episode limitation
STEP = 300 # Step limitation in an episode
TEST = 10 # The number of experiment test every 100 episode

def main():
	# initialize OpenAI Gym env and dqn agent
	env = gym.make(ENV_NAME)
	agent = DQN(env)

	for episode in range(EPISODE):
		# initialize task
		state_narray = env.reset()
		state = state_narray.flatten()
		# Train 
		for step in range(STEP):
    		# problem
			action = agent.egreedy_action(state) # e-greedy action for train  # 获取包含随机的动作,输出动作
			next_state_narray,reward,done,_ = env.step(action) # 然后环境自己执行动作，输出新的状态
			next_state = next_state_narray.flatten()
			# Define reward for agent
			# reward_agent = -1 if done else 0.1
			agent.perceive(state,action,reward,next_state,done) # 感知信息,输入状态
			state = next_state
			if done:
				break
		# Test every 100 episodes
		if episode % 100 == 0:
			total_reward = 0
			for i in range(TEST):
				state_narray = env.reset()
				state = state_narray.flatten()
				for j in range(STEP):
					env.render()
					# problem
					action = agent.action(state) # direct action for test,测试中唯一的不同,来获取动作，也就是完全没有随机性，只根据神经网络来输出，没有探索，同时这里也就不再perceive输入信息来训练
					state_narray,reward,done,_ = env.step(action)
					state = state_narray.flatten()
					total_reward += reward
					if done:
						break
			ave_reward = total_reward/TEST
			print ('episode: ',episode,'Evaluation Average Reward:',ave_reward)
			if ave_reward >= 200:
				break

	# # save results for uploading
	# env.monitor.start('gym_results/CartPole-v0-experiment-1',force = True)
	# for i in range(100):
	# 	state = env.reset()
	# 	for j in range(200):
	# 		env.render()
	# 		action = agent.action(state) # direct action for test
	# 		state,reward,done,_ = env.step(action)
	# 		total_reward += reward
	# 		if done:
	# 			break
	# env.monitor.close()

if __name__ == '__main__':
	main()